package CPU.rv64_5stage

import chisel3._
import chisel3.util._

class EX2MEM extends Bundle{
  val exout   = UInt(64.W)
  val rfwaddr = UInt(5.W)
  val rfwen   = Bool()
  val opcode  = UInt(7.W)
  val islsu   = Bool()
  val pcinst  = new IF2ID

  val en    = Bool()
  val addr  = UInt(64.W)
  val wdata  = UInt(64.W)
  val wmask  = UInt(64.W)
  val wen    = Bool()
}


class Execution extends Module{
  val io = IO(new Bundle{
    val exin    = Input(new ID2EX)
    val ex2mem  = Output(new EX2MEM)

  })
  io.ex2mem.pcinst  := io.exin.pcinst
  io.ex2mem.rfwaddr := io.exin.rfwaddr
  io.ex2mem.rfwen   := io.exin.rfwen
  io.ex2mem.opcode  := io.exin.opcode
  io.ex2mem.islsu   := io.exin.islsu

  val alu_out = MuxCase(0.U,Array(
    (io.exin.opcode===ALUOpType.add)->(io.exin.in1 + io.exin.in2) ,
    (io.exin.opcode===ALUOpType.sub)->(io.exin.in1 - io.exin.in2) ,
    (io.exin.opcode===ALUOpType.and)->(io.exin.in1 & io.exin.in2) ,
    (io.exin.opcode===ALUOpType.or )->(io.exin.in1 | io.exin.in2) ,
    (io.exin.opcode===ALUOpType.xor)->(io.exin.in1 ^ io.exin.in2) ,
    (io.exin.opcode===ALUOpType.slt)->(io.exin.in1.asSInt() < io.exin.in2.asSInt()).asUInt() ,
    (io.exin.opcode===ALUOpType.sltu)->(io.exin.in1 < io.exin.in2).asUInt() ,
    (io.exin.opcode===ALUOpType.sll)->(io.exin.in1 << io.exin.in2(5,0)),
    (io.exin.opcode===ALUOpType.sra)->(io.exin.in1.asSInt() >> io.exin.in2(5,0)).asUInt() ,
    (io.exin.opcode===ALUOpType.srl)->(io.exin.in1 >> io.exin.in2(5,0)),

    (io.exin.opcode===ALUOpType.addw)->SignExt((io.exin.in1 +  io.exin.in2)(31,0),64) ,
    (io.exin.opcode===ALUOpType.subw)->SignExt((io.exin.in1 -  io.exin.in2)(31,0),64) ,
    (io.exin.opcode===ALUOpType.sllw)->SignExt((io.exin.in1 << io.exin.in2(4,0))(31,0),64) ,
    (io.exin.opcode===ALUOpType.srlw)->SignExt((io.exin.in1(31,0) >> io.exin.in2(4,0)).asUInt(),64) ,
    (io.exin.opcode===ALUOpType.sraw)->SignExt((io.exin.in1(31,0).asSInt() >> io.exin.in2(4,0)).asUInt(),64),
    (io.exin.opcode===ALUOpType.jal)->(io.exin.in1 + io.exin.in2) ,
  ))
  io.ex2mem.exout := Mux(io.exin.isjalr,ZeroExt(io.exin.pcinst.pc+4.U,64),alu_out.asUInt())

  val addroffset = WireInit(0.U(3.W))
  val offsetdata = WireInit(0.U(64.W))

  io.ex2mem.en    := io.exin.islsu
  io.ex2mem.addr  := io.exin.in1 + io.exin.in2
  addroffset      := io.ex2mem.addr(2,0)
  io.ex2mem.wen   := (io.exin.islsu) && LSUOpType.isStore(io.exin.opcode)
  io.ex2mem.wdata := io.exin.rs2 << (addroffset<<3)

  val mask = MuxCase(0.U,Array(
    (io.exin.opcode===LSUOpType.sb) -> ZeroExt("hff".U,64),
    (io.exin.opcode===LSUOpType.sh) -> ZeroExt("hffff".U,64),
    (io.exin.opcode===LSUOpType.sw) -> ZeroExt("hffffffff".U,64),
    (io.exin.opcode===LSUOpType.sd) -> (-1).S(64.W).asUInt(),
  ))
  io.ex2mem.wmask := mask << (addroffset<<3)
}


object SignExt {
  def apply(a: UInt, len: Int) = {
    val aLen = a.getWidth
    val signBit = a(aLen-1)
    if (aLen >= len) a(len-1,0) else Cat(Fill(len - aLen, signBit), a)
  }
}

object ZeroExt {
  def apply(a: UInt, len: Int) = {
    val aLen = a.getWidth
    if (aLen >= len) a(len-1,0) else Cat(0.U((len - aLen).W), a)
  }
}